import colorsys
import copy
import time

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn

from riversr.model.u_net.nets.unet import Unet as unet
from riversr.model.u_net.utils.utils import cvtColor, preprocess_input, resize_image, show_config

#--------------------------------------------#
#   使用自己训练好的模型预测需要修改2个参数
#   model_path和num_classes都需要修改！
#   如果出现shape不匹配
#   一定要注意训练时的model_path和num_classes数的修改
#--------------------------------------------#
class Unet(object):
    _defaults = {
        #-------------------------------------------------------------------#
        #   model_path指向logs文件夹下的权值文件
        #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
        #   验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。
        #-------------------------------------------------------------------#
        # "model_path"    : 'model_data/unet_resnet_voc.pth',
        "model_path"    : r'riversr\model\u_net\logs\best_epoch_weights.pth',
        #--------------------------------#
        #   所需要区分的类的个数+1
        #--------------------------------#
        "num_classes"   : 2,
        #--------------------------------#
        #   所使用的的主干网络：vgg、resnet50
        #--------------------------------#
        "backbone"      : "resnet50",
        #--------------------------------#
        #   输入图片的大小
        #--------------------------------#
        "input_shape"   : [512, 512],
        #-------------------------------------------------#
        #   mix_type参数用于控制检测结果的可视化方式
        #
        #   mix_type = 0的时候代表原图与生成的图进行混合          具体表现：原图变黑，目标红色
        #   mix_type = 1的时候代表仅保留生成的图                 具体表现：目标红色，其余黑的
        #   mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标  具体表现：目标分割，其余黑的
        #-------------------------------------------------#
        "mix_type"      : 0,
        #--------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        #--------------------------------#
        "cuda"          : True,
    }

    #---------------------------------------------------#
    #   初始化UNET
    #---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        for name, value in kwargs.items():
            setattr(self, name, value)
        #---------------------------------------------------#
        #   画框设置不同的颜色
        #---------------------------------------------------#
        if self.num_classes <= 21:
            self.colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128), 
                            (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128), 
                            (64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128), 
                            (128, 64, 12)]
        else:
            hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
            self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
            self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
        #---------------------------------------------------#
        #   获得模型
        #---------------------------------------------------#
        self.generate()
        
        show_config(**self._defaults)

    #---------------------------------------------------#
    #   获得所有的分类
    #---------------------------------------------------#
    def generate(self, onnx=False):
        self.net = unet(num_classes = self.num_classes, backbone=self.backbone)

        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
        self.net    = self.net.eval()
        print('{} model, and classes loaded.'.format(self.model_path))
        if not onnx:
            if self.cuda:
                self.net = nn.DataParallel(self.net)
                self.net = self.net.cuda()

    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def detect_image_pro(self, image, count=False, name_classes=None):
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------#
        #   对输入图像进行一个备份，后面用于绘图
        #---------------------------------------------------#
        old_img     = copy.deepcopy(image)
        orininal_h  = np.array(image).shape[0]
        orininal_w  = np.array(image).shape[1]
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data, nw, nh  = resize_image(image, (self.input_shape[1],self.input_shape[0]))
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
                
            #---------------------------------------------------#
            #   图片传入网络进行预测
            #---------------------------------------------------#
            pr = self.net(images)[0]
            #---------------------------------------------------#
            #   取出每一个像素点的种类
            #---------------------------------------------------#
            pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy()
            #--------------------------------------#
            #   将灰条部分截取掉
            #--------------------------------------#
            pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
                    int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
            #---------------------------------------------------#
            #   进行图片的resize
            #---------------------------------------------------#
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)
            #---------------------------------------------------#
            #   取出每一个像素点的种类
            #---------------------------------------------------#
            pr = pr.argmax(axis=-1)
        
        if self.mix_type == 0:
            # seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))
            # for c in range(self.num_classes):
            #     seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')
            #     seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')
            #     seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            #------------------------------------------------#
            #   将新图片转换成Image的形式
            #------------------------------------------------#
            image   = Image.fromarray(np.uint8(seg_img))
            # 计算image中目标的宽高
            # 计算高
            
            # 计算宽 图片中目标的宽度，目标是红色的(128,0,0),背景是黑色的(0,0,0)，获取最左端和最右端的像素点的位置
            # 获取最左端的像素点,按列-行的顺序获取

            left = (0,0)
            for i in range(orininal_w//2,orininal_w):
                for j in range(orininal_h//2,orininal_h):
                    if image.getpixel((i,j)) == (128,0,0):
                        left = (i,j)
                        break
                if left != (0,0):
                    break
                    
            right = (0,0)
            for i in range(orininal_w-1,-1,-1):
                for j in range(orininal_h//2,orininal_h):
                    if image.getpixel((i,j)) == (128,0,0):
                        right = (i,j)
                        break
                if right != (0,0):
                    break
            # 计算欧几里得距离
            distance = np.sqrt((right[0]-left[0])**2 + (right[1]-left[1])**2)

            # 可知220px = 0.4m
            level = 727.82 - distance * 0.4 / 220
            # 四舍五入
            level = round(level,2)

            # 在两点之间画线 PIL
            from PIL import ImageDraw
            # image = Image.fromarray(np.uint8(seg_img))
            old_img = Image.fromarray(np.uint8(old_img))
            # draw = ImageDraw.Draw(image)
            draw_oldimg = ImageDraw.Draw(old_img)
            # draw.line([left, right], fill=(0, 255, 0), width=5)

            # y=kx+b ,过点 left 画一条线，k=1, b = left[1] - left[0]
            point_x = (1450, 0)
            point_y = (1350, orininal_h)
            # 根据两点，计算k
            k = (point_y[1] - point_x[1]) / (point_y[0] - point_x[0])
            # 根据 left 计算b
            b = left[1] - k * left[0]

            # 重新调整point_x的位置和point_y的位置，使其过left点，x点在图片顶部，y点在图片底部
            point_x = (int((0 - b) / k), 0)
            point_y = (int((orininal_h - b) / k), orininal_h)

            
            # draw.line([point_x, point_y], fill=(0, 255, 0), width=5)
            draw_oldimg.line([point_x, point_y], fill=(0, 255, 0), width=5)
            # 在图片上写字
            from PIL import ImageFont
            font = ImageFont.truetype("arial.ttf", 100)
            draw_oldimg.text((300, 500), "Water Level: {} m".format(level), fill=(255, 0, 0), font=font)
            
            print("left:",left,"right:",right,"distance:",distance,"level:",level,"k:",k)
            #------------------------------------------------#
            #   将新图与原图及进行混合
            #------------------------------------------------#
            # image   = Image.blend(old_img, image, 0.7)
            image   = old_img

        elif self.mix_type == 1:
            # seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))
            # for c in range(self.num_classes):
            #     seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')
            #     seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')
            #     seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            #------------------------------------------------#
            #   将新图片转换成Image的形式
            #------------------------------------------------#
            image   = Image.fromarray(np.uint8(seg_img))

        elif self.mix_type == 2:
            seg_img = (np.expand_dims(pr != 0, -1) * np.array(old_img, np.float32)).astype('uint8')
            #------------------------------------------------#
            #   将新图片转换成Image的形式
            #------------------------------------------------#
            image = Image.fromarray(np.uint8(seg_img))
        
        return image
    

    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def detect_image_pro_str(self, image, count=False, name_classes=None):
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------#
        #   对输入图像进行一个备份，后面用于绘图
        #---------------------------------------------------#
        old_img     = copy.deepcopy(image)
        orininal_h  = np.array(image).shape[0]
        orininal_w  = np.array(image).shape[1]
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data, nw, nh  = resize_image(image, (self.input_shape[1],self.input_shape[0]))
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
                
            #---------------------------------------------------#
            #   图片传入网络进行预测
            #---------------------------------------------------#
            pr = self.net(images)[0]
            #---------------------------------------------------#
            #   取出每一个像素点的种类
            #---------------------------------------------------#
            pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy()
            #--------------------------------------#
            #   将灰条部分截取掉
            #--------------------------------------#
            pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
                    int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
            #---------------------------------------------------#
            #   进行图片的resize
            #---------------------------------------------------#
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)
            #---------------------------------------------------#
            #   取出每一个像素点的种类
            #---------------------------------------------------#
            pr = pr.argmax(axis=-1)
        
        level = 0
        if self.mix_type == 0:
            # seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))
            # for c in range(self.num_classes):
            #     seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')
            #     seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')
            #     seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')
            seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
            #------------------------------------------------#
            #   将新图片转换成Image的形式
            #------------------------------------------------#
            image   = Image.fromarray(np.uint8(seg_img))
            # 计算image中目标的宽高
            # 计算高
            
            # 计算宽 图片中目标的宽度，目标是红色的(128,0,0),背景是黑色的(0,0,0)，获取最左端和最右端的像素点的位置
            # 获取最左端的像素点,按列-行的顺序获取

            left = (0,0)
            for i in range(orininal_w//2,orininal_w):
                for j in range(orininal_h//2,orininal_h):
                    if image.getpixel((i,j)) == (128,0,0):
                        left = (i,j)
                        break
                if left != (0,0):
                    break
                    
            right = (0,0)
            for i in range(orininal_w-1,-1,-1):
                for j in range(orininal_h//2,orininal_h):
                    if image.getpixel((i,j)) == (128,0,0):
                        right = (i,j)
                        break
                if right != (0,0):
                    break
            # 计算欧几里得距离
            distance = np.sqrt((right[0]-left[0])**2 + (right[1]-left[1])**2)

            # 可知220px = 0.4m
            level = 727.82 - distance * 0.4 / 220
            # 四舍五入
            level = round(level,2)

            return level
        
        return image